import backtrader as bt


class SSAverange(bt.Strategy):
    """
    继承并构建自己的bt策略
    """

    def log(self, txt, dt=None, doprint=True):
        ''' 日志函数，用于统一输出日志格式 '''
        if doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # 初始化相关数据
        self.dataclose = self.datas[0].close

        self.order = None
        self.buyprice = None
        self.buycomm = None

        # 五分钟移动平均线
        self.sma5 = bt.indicators.MovingAverageSimple(self.datas[0], period=2)
        # 十分钟移动平均线
        self.sma10 = bt.indicators.MovingAverageSimple(self.datas[0], period=17)
        self.buy_signal = bt.indicators.CrossOver(self.sma5, self.sma10)  # 设置买入信号
        self.sell_signal = bt.indicators.CrossDown(self.sma5, self.sma10)  # 设置卖出信号
        # rsi 超过80说明将要下降
        # rsi 小于20说明将要上升

        self.buy_signal.plotinfo.plot = False
        self.sell_signal.plotinfo.plot = False

    def notify_order(self, order):
        """
        在next方法调用之前调用
        订单中任意状态改变 ? 需要再研究
        订单状态处理

        Arguments:
            order {object} -- 订单状态
        """
        if order.status in [order.Submitted, order.Accepted]:
            # 如订单已被处理，则不用做任何事情
            return

        # 检查订单是否完成
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.log("it's buy order time is %s the price is %s " %
                      (self.datas[0].datetime.datetime(0), self.buyprice), doprint=True)
            if order.issell():
                self.sellprice = order.executed.price
                self.sellcomm = order.executed.comm
                self.log("it's sell order time is %s the price is %s " %
                      (self.datas[0].datetime.datetime(0), self.sellprice), doprint=True)
            self.bar_executed = len(self)

        # 订单因为缺少资金之类的原因被拒绝执行
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected', doprint=True)

        # 订单状态处理完成，设为空
        self.order = None

    def notify_trade(self, trade):
        """
        开仓，更新，平仓
        交易成果

        Arguments:
            trade {object} -- 交易状态
        """
        if not trade.isclosed:
            return

        # 显示交易的毛利率和净利润
        self.log('OPERATION PROFIT, 毛利率 %.2f, 净利润 %.2f' %
                 (trade.pnl, trade.pnlcomm), doprint=False)

    def next(self):
        # 记录收盘价
        self.log('Close, %.2f' % self.dataclose[0], doprint=False)
        # 是否正在下单，如果是的话不能提交第二次订单
        if self.order:
            return

        # 是否已经买入
        if not self.position and self.buy_signal[0] == 1:
            self.order = self.buy()
            self.buy_flag = 0

        # if self.sell_flag == 1:
        if self.position and self.sell_signal[0] == 1:
            self.order = self.sell()
            self.sell_flag = 0

    def stop(self):
        self.log(u'(金叉死叉有用吗) Ending Value %.2f' %
                 (self.broker.getvalue()), doprint=True)
